package com.dl4jlpr;

import lombok.extern.slf4j.Slf4j;
import org.datavec.api.io.filters.RandomPathFilter;
import org.datavec.api.split.FileSplit;
import org.datavec.api.split.InputSplit;
import org.datavec.image.loader.NativeImageLoader;
import org.datavec.image.recordreader.ImageRecordReader;
import org.deeplearning4j.datasets.datavec.RecordReaderDataSetIterator;
import org.deeplearning4j.datasets.iterator.INDArrayDataSetIterator;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.*;
import org.deeplearning4j.nn.conf.graph.LayerVertex;
import org.deeplearning4j.nn.conf.graph.StackVertex;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.common.io.ClassPathResource;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.learning.config.Nesterovs;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.io.File;
import java.io.IOException;
import java.util.*;


@Slf4j
public class LprNet {


    public static void detection() throws Exception {
        int width = 94;
        int height = 24;
        int nChannels = 3;
        int batchSize = 8;
        int nEpochs = 10;
        double learningRate = 1e-3;
        double lrMomentum = 0.9;
        int seed = 1000;
        Random rng = new Random(seed);
        RandomPathFilter pathFilter = new RandomPathFilter(rng);
        String modelFilename = "LprNetModel.zip";
        String dataDir = new ClassPathResource("/").getFile().getPath();
        File imageDir = new File(dataDir, "test");
        log.info("Load LPRNet train data...");
        //log.info(Arrays.toString(PlateLableGenerator.CHARS));
        InputSplit[] data = new FileSplit(imageDir, NativeImageLoader.ALLOWED_FORMATS, rng).sample(pathFilter, 0.8, 0.2);
        InputSplit trainData = data[0];

        InputSplit testData = data[1];

        LprIterable lpr=new LprIterable(24,94,3);
        INDArrayDataSetIterator train = new INDArrayDataSetIterator(lpr.CreatIterable(data[0]),1);

                ComputationGraphConfiguration conf =
                new NeuralNetConfiguration.Builder().seed(seed)
                        .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT).gradientNormalization(GradientNormalization.RenormalizeL2PerLayer)
                        .gradientNormalizationThreshold(1.0).updater(new Adam.Builder().learningRate(learningRate).build())
                        .updater(new Nesterovs.Builder().learningRate(learningRate).momentum(lrMomentum).build())
                        .activation(Activation.RELU)
                        .trainingWorkspaceMode(WorkspaceMode.ENABLED).inferenceWorkspaceMode(WorkspaceMode.ENABLED)
                        .graphBuilder()
                        .addInputs("in")
                        .layer("con1",new ConvolutionLayer.Builder(1, 1).nIn(3).nOut(1).stride(1, 1).convolutionMode(ConvolutionMode.Same)
                                .weightInit(WeightInit.UNIFORM).hasBias(false).activation(Activation.IDENTITY).build(),"in")
                        .layer("test", new SubsamplingLayer.Builder().poolingType(SubsamplingLayer.PoolingType.MAX).kernelSize(24,3)
                                .stride(14, 14).build(),"con1")
                        .addVertex("test1",new RiseJoin(8,1),"test")
                        .layer("out",new OutputLayer.Builder().nIn(7).nOut(7).activation(Activation.SOFTMAX)
                                .weightInit(WeightInit.XAVIER)
                                .lossFunction(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD).build(),"test1")
                        .setOutputs("out")
                        .setInputTypes(InputType.convolutionalFlat(24,94,3))
                        .build();


        ComputationGraph model = new ComputationGraph(conf);
        model.init();

        log.info("LPRNET is:{}",model.summary(InputType.convolutional(height, width, nChannels)));
        log.info("Train model...");

        model.setListeners(new ScoreIterationListener(1));
        for (int i = 0; i < nEpochs; i++) {
            train.reset();
            while (train.hasNext()) {
                //DataSet ds = train.next();
                //System.out.println(ds);
                //log.info(ds);
                model.fit(train.next());
            }
            log.info("*** Completed epoch {} ***", i);
        }
        ModelSerializer.writeModel(model, modelFilename, true);
        //return network;
    }
    public static void main(String[] args) throws Exception {
        detection();
    }
}
/*
//.addVertex("debug",new RiseJoin(),"in")
//.layer("conv1", new Convolution2D.Builder().kernelSize(1, 1).stride(1,1).nIn(3).nOut(3).build(),"in")
//.addVertex("debug",new RiseJoin(),"conv1")

                        .layer("conv1", new Convolution2D.Builder().kernelSize(3, 3).stride(1,1).nIn(3).nOut(64).build(),"in")
                                .layer("BN2d_1",new BatchNormalization.Builder().nOut(64).build(),"conv1")
                                .layer("active_1", new ActivationLayer.Builder().activation(Activation.RELU).build(), "BN2d_1")
                                //.layer("MaxPool3d", new Subsampling3DLayer.Builder(Subsampling3DLayer.PoolingType.MAX).kernelSize(1,3, 3).stride(1, 1,1)
                                //       .build(),"active_1")
                                .layer("MP3d_ts0",new MaxPool_Fake3d.Builder().dimensions1(0,2,3,1).dimensions2(0,3,1,2).kernelSize(1,3,3).strideSize(1,1,1).nIn(64).nOut(64).build(),"active_1")
                                //.layer("MP2d_1", new SubsamplingLayer.Builder(PoolingType.MAX).kernelSize(3, 3).stride( 1,1).build(),"active_1")
                                //.layer("MP3d_1", new Subsampling3DLayer.Builder().kernelSize(1, 3,3).poolingType(Subsampling3DLayer.PoolingType.MAX).stride(1, 1,1).build(),"active_1")
                                //.layer("sbb_11", new ConvolutionLayer.Builder().kernelSize(1, 1).nIn(64).nOut(128).build(),"MP2d_1")
                                //.layer("sbb_12", new ConvolutionLayer.Builder().kernelSize(3, 1).padding(1,0).nIn(128).nOut(128).build(),"sbb_11")
                                //.layer("sbb_13", new ConvolutionLayer.Builder().kernelSize(1, 3).padding(0,1).nIn(128).nOut(128).build(),"sbb_12")
                                //.layer("sbb_14", new ConvolutionLayer.Builder().kernelSize(1, 1).nIn(128).nOut(128).build(),"sbb_13")
                                .layer("sbb_1",new small_basic_block.Builder().nIn(64).nOut(128).batchsize(8).build(),"MP3d_ts0")

                                .layer("BN2d_2",new BatchNormalization.Builder().nIn(128).nOut(128).build(),"sbb_1")
                                .layer("active_2", new ActivationLayer.Builder().activation(Activation.RELU).build(), "BN2d_2")

                                .layer("MP3d_ts1",new MaxPool_Fake3d.Builder().dimensions1(0,2,3,1).dimensions2(0,3,1,2).kernelSize(1,3,3).strideSize(2,1,2).nIn(128).nOut(64).build(),"active_2")

                                //.layer("sbb_21", new ConvolutionLayer.Builder().kernelSize(1, 1).nIn(64).nOut(256).build(),"MP3d_ts1")
                                //.layer("sbb_22", new ConvolutionLayer.Builder().kernelSize(3, 1).padding(1,0).nIn(256).nOut(256).build(),"sbb_21")
                                //.layer("sbb_23", new ConvolutionLayer.Builder().kernelSize(1, 3).padding(0,1).nIn(256).nOut(256).build(),"sbb_22")
                                //.layer("sbb_24", new ConvolutionLayer.Builder().kernelSize(1, 1).nIn(256).nOut(256).build(),"sbb_23")
                                .layer("sbb_2",new small_basic_block.Builder().nIn(64).nOut(256).batchsize(8).build(),"MP3d_ts1")

                                .layer("BN2d_3",new BatchNormalization.Builder().nOut(256).build(),"sbb_2")
                                .layer("active_3", new ActivationLayer.Builder().activation(Activation.RELU).build(), "BN2d_3")

                                //.layer("sbb_31", new ConvolutionLayer.Builder().kernelSize(1, 1).nIn(256).nOut(256).build(),"active_3")
                                //.layer("sbb_32", new ConvolutionLayer.Builder().kernelSize(3, 1).padding(1,0).nIn(256).nOut(256).build(),"sbb_31")
                                //.layer("sbb_33", new ConvolutionLayer.Builder().kernelSize(1, 3).padding(0,1).nIn(256).nOut(256).build(),"sbb_32")
                                //.layer("sbb_34", new ConvolutionLayer.Builder().kernelSize(1, 1).nIn(256).nOut(256).build(),"sbb_33")
                                .layer("sbb_3",new small_basic_block.Builder().nIn(256).nOut(256).batchsize(8).build(),"active_3")

                                .layer("BN2d_4",new BatchNormalization.Builder().nOut(256).build(),"sbb_3")
                                .layer("active_4", new ActivationLayer.Builder().activation(Activation.RELU).build(), "BN2d_4")

                                .layer("MP3d_ts2",new MaxPool_Fake3d.Builder().dimensions1(0,2,3,1).dimensions2(0,3,1,2).kernelSize(1,3,3).strideSize(4,1,2).nIn(256).nOut(64).build(),"active_4")


                                .layer("d1",new DropoutLayer.Builder(0.5).build(),"MP3d_ts2")

                                .layer("conv2", new ConvolutionLayer.Builder().kernelSize(1, 4).stride(1,1).nIn(64).nOut(256).build(),"d1")
                                .layer("BN2d_5",new BatchNormalization.Builder().nOut(256).build(),"conv2")
                                .layer("active_5", new ActivationLayer.Builder().activation(Activation.RELU).build(), "BN2d_5")

                                .layer("d2",new DropoutLayer.Builder(0.5).build(),"active_5")

                                .layer("conv3", new ConvolutionLayer.Builder().kernelSize(13, 1).stride(1,1).nIn(256).nOut(68).build(),"d2")
                                .layer("BN2d_6",new BatchNormalization.Builder().nOut(68).build(),"conv3")
                                .layer("active_6", new ActivationLayer.Builder().activation(Activation.RELU).build(), "BN2d_6")
                                //.layer("conv4", new ConvolutionLayer.Builder().kernelSize(1, 1).stride(1,1).nIn(516).nOut(68).build(),"active_6")

                                .addLayer("output",  new LossLayer.Builder().build(), "active_6")
                                .setOutputs("output")
//.layer("output", new OutputLayer.Builder(
//        LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD).nIn(66).nOut(66)
//        .activation(Activation.SOFTMAX) // radial basis function required
//        .build(), "conv4")
//.setOutputs("output")

 */
